package com.minio.config;

import com.alibaba.fastjson.JSONObject;
import com.google.gson.*;
import lombok.extern.slf4j.Slf4j;
import okhttp3.*;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.web.multipart.MultipartFile;

import javax.annotation.PostConstruct;
import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.*;
import java.net.URLEncoder;
import java.nio.charset.StandardCharsets;
import java.time.Duration;
import java.time.Instant;
import java.util.concurrent.TimeUnit;

@Slf4j
@Component
public class ComfyUIImg2ImgWithOkHttp {
    private OkHttpClient client;
    
    @Autowired
    private ComfyUIProperties comfyUIProperties;
    
    private String comfyBaseUrl;

    /**
     * 初始化客户端
     */
    @PostConstruct
    public void init() {
        this.comfyBaseUrl = comfyUIProperties.getBaseUrl().endsWith("/")
                ? comfyUIProperties.getBaseUrl().substring(0, comfyUIProperties.getBaseUrl().length() - 1)
                : comfyUIProperties.getBaseUrl();

        // 配置OkHttpClient
        this.client = new OkHttpClient.Builder()
                .connectTimeout(comfyUIProperties.getTimeout(), TimeUnit.SECONDS)
                .readTimeout(comfyUIProperties.getTimeout(), TimeUnit.SECONDS)
                .writeTimeout(comfyUIProperties.getTimeout(), TimeUnit.SECONDS)
                .build();
        boolean isConnected = testConnection();
        if (isConnected) {
            System.out.println("ComfyUI 连接成功");
        } else {
            System.out.println("ComfyUI 连接失败");
        }
    }

    /**
     * 测试与ComfyUI的连接
     */
    public boolean testConnection() {
        System.out.println(comfyBaseUrl);
        Request request = new Request.Builder()
                .url(comfyBaseUrl + "/system_stats")
                .get()
                .build();

        try (Response response = client.newCall(request).execute()) {
            return response.isSuccessful();
        } catch (IOException e) {
            System.err.println("连接测试失败: " + e.getMessage());
            return false;
        }
    }

    /**
     * 上传图片到ComfyUI（使用Base64方式，兼容官方API）
     * @param file
     * @return 上传后的图片标识名
     */
    public String uploadFile(File file,String fileName) {
        try {

            // 构建multipart请求体
            MultipartBody requestBody = new MultipartBody.Builder()
                    .setType(MultipartBody.FORM)
                    // 图片文件字段（关键：字段名需与ComfyUI接口期望一致）
                    .addFormDataPart("image", fileName,RequestBody.create(MediaType.parse("image/*"), file))
                    // 额外参数：指定上传目录（ComfyUI通常需要）
                    .addFormDataPart("type", "input")
                    .build();

            // 发送请求
            Request request = new Request.Builder()
                    .url(comfyBaseUrl + "/upload/image")
                    .post(requestBody)
                    .build();

            try (Response response = client.newCall(request).execute()) {
                if (!response.isSuccessful()) {
                    throw new IOException("上传失败: " + response.code() + " " + response.message());
                }
                if (response.body() == null) {
                    return null;
                }
                // 解析响应
                String responseBody = response.body().string();
                JsonObject jsonResponse = JsonParser.parseString(responseBody).getAsJsonObject();
                String name = jsonResponse.get("name").getAsString();

                return name;
            }
        } catch (Exception e) {
            System.err.println("图片上传失败: " + e.getMessage());
            return null;
        }
    }

    /**
     * 提交图生图任务
     * @param json 任务配置参数
     * @return 任务ID
     */
    public String submitImg2ImgTask(String json) {
        try {
            RequestBody body = RequestBody.create(MediaType.parse("application/json; charset=utf-8"), json);
            // 发送请求
            Request request = new Request.Builder()
                    .url(comfyBaseUrl + "/prompt")
                    .post(body)
                    .build();

            try (Response response = client.newCall(request).execute()) {
                if (!response.isSuccessful()) {
                    throw new IOException("提交任务失败: " + response.code());
                }
                if (response.body() == null) {
                    return null;
                }
                String responseBody = response.body().string();
                JsonObject jsonResponse = JsonParser.parseString(responseBody).getAsJsonObject();
                return jsonResponse.get("prompt_id").getAsString();
            }
        } catch (Exception e) {
            System.err.println("提交任务失败: " + e.getMessage());
            return null;
        }
    }

    /**
     * 查询任务状态
     * @param taskId 任务ID
     * @return 任务状态JSON对象
     */
    public JsonObject getTaskStatus(String taskId) {
        try {
            Request request = new Request.Builder()
                    .url(comfyBaseUrl + "/history/" + taskId)
                    .get()
                    .build();

            try (Response response = client.newCall(request).execute()) {
                if (response.isSuccessful() && response.body() != null) {
                    String responseBody = response.body().string();
                    return JsonParser.parseString(responseBody).getAsJsonObject();
                }
            }
        } catch (Exception e) {
            System.err.println("查询任务状态失败: " + e.getMessage());
        }
        return null;
    }

    /**
     * 等待任务完成
     * @param taskId 任务ID
     * @param checkInterval 检查间隔(毫秒)
     * @param maxWaitTime 最大等待时间(毫秒)
     * @return 任务成功完成返回true
     */
    public boolean waitForTaskCompletion(String taskId, int checkInterval, long maxWaitTime) {
        Instant start = Instant.now();
        System.out.println("等待任务 [" + taskId + "] 完成...");

        while (true) {
            // 检查超时
            if (Duration.between(start, Instant.now()).toMillis() > maxWaitTime) {
                System.err.println("任务超时");
                return false;
            }
            // 查询任务状态
            JsonObject status = getTaskStatus(taskId);
            if (status != null && status.has(taskId)) {
                JsonObject taskInfo = status.getAsJsonObject(taskId);
                if (taskInfo.has("error")) {
                    System.err.println("任务执行错误: " + taskInfo.get("error").getAsString());
                    return false;
                }
                System.out.println("任务完成!");
                return true;
            }

            // 等待下次检查
            try {
                Thread.sleep(checkInterval);
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
                return false;
            }
        }
    }

    /**
     * 获取生成的图片
     * @param filename 图片文件名
     * @param subfolder 子文件夹
     * @param type 文件夹类型
     * @return 图片对象
     */
    public BufferedImage getGeneratedImage(String filename, String subfolder, String type) {
        try {
            // 编码URL参数
            String encodedFilename = URLEncoder.encode(filename, StandardCharsets.UTF_8.name());
            String encodedSubfolder = URLEncoder.encode(subfolder, StandardCharsets.UTF_8.name());
            String encodedType = URLEncoder.encode(type, StandardCharsets.UTF_8.name());

            // 构建图片URL
            String url = comfyBaseUrl + "/view?filename=" + encodedFilename +
                    "&subfolder=" + encodedSubfolder + "&type=" + encodedType;

            Request request = new Request.Builder()
                    .url(url)
                    .get()
                    .build();

            try (Response response = client.newCall(request).execute()) {
                if (response.isSuccessful() && response.body() != null) {
                    return ImageIO.read(response.body().byteStream());
                } else {
                    throw new IOException("获取图片失败，响应码: " + response.code());
                }
            }
        } catch (Exception e) {
            System.err.println("获取生成图片失败: " + e.getMessage());
            return null;
        }
    }


}